
import apriltag   
import cv2
import numpy as np

class AprilTag():
    def __init__(self):
        self.detector = apriltag.Detector(apriltag.DetectorOptions(families='tag36h11')) # windows

    def detect(self,img):
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        tags = self.detector.detect(gray)
        returnList = []
        self.returnListFourPoint = []
        for tag in tags:
            point = np.array(tuple(i.astype(int) for i in tag.corners))
            self.returnListFourPoint.append(point)
            [cv2.circle(img, tuple(p), 4, (255, 0, 255), 2) for p in point]

            centerpoint = [int(np.mean(point[:,0])),int(np.mean(point[:,1]))]
            centerpointFloat = [np.mean(point[:,0]),np.mean(point[:,1])]
            cv2.circle(img, tuple(centerpoint), 4, (0, 255, 0), 10)

            returnList.append(centerpointFloat)

        return returnList,img

    def PNP(self,tag_size,tag_points):
        from Main.Class.Message import MTX,DIST
        
        tag_3d_points = np.array([[0, 0, 0], [tag_size, 0, 0], [tag_size, tag_size, 0], [0, tag_size, 0]], dtype=np.float32)

        success, rvec, tvec = cv2.solvePnP(tag_3d_points, np.asarray(tag_points,dtype="float32"), MTX, DIST)

        return rvec, tvec

def rotation_matrix_to_euler_angles(R):
    sy = np.sqrt(R[0, 0] * R[0, 0] + R[1, 0] * R[1, 0])

    singular = sy < 1e-6

    if not singular:
        x = np.arctan2(R[2, 1], R[2, 2])
        y = np.arctan2(-R[2, 0], sy)
        z = np.arctan2(R[1, 0], R[0, 0])
    else:
        x = np.arctan2(-R[1, 2], R[1, 1])
        y = np.arctan2(-R[2, 0], sy)
        z = 0

    return np.degrees(x), np.degrees(y), np.degrees(z)

if __name__ == "__main__":
    a = AprilTag()

    img = cv2.imread("/home/shansu/work/ws_panda/ws_panda/1.jpg")
    center,img = a.detect(img)
    print(center) #[[509.5, 267.5]]
    print(a.returnListFourPoint)
    rvec,tvec = a.PNP(160,a.returnListFourPoint[0])

    rmtx, _ = cv2.Rodrigues(rvec)

    print(rmtx)
    print("roll yaw pitch")
    print(rotation_matrix_to_euler_angles(rmtx))


    cv2.imshow("out_image", img)
    cv2.imwrite('image.png', img)
    cv2.waitKey()
